#!/usr/bin/env python3

import sys
import re
from datetime import datetime
from typing import Optional, Tuple, Dict
from collections import namedtuple
from elftools.elf.elffile import ELFFile
from elftools.dwarf.compileunit import CompileUnit
from elftools.dwarf.die import DIE, AttributeValue
from elftools.dwarf.descriptions import ExprDumper, set_global_machine_arch

debug_mode = 0
cache_enabled = False

if debug_mode:
    def log(*args, **kwargs): print(*args, **kwargs)
    def die_name(die: DIE):
        if die is None: return None
        try:
            return die.attributes.get('DW_AT_name').value.decode('utf-8')
        except AttributeError:
            return '<DieNameErr>'
else:
    def log(*args, **kwargs): pass
    def die_name(die: DIE): pass


class LocTuple(namedtuple('LocTuple', ['register', 'offset'])):
    @staticmethod
    def from_str(line):
        m = re.search(r'DW_OP_fbreg: (-?\d+)', line)
        if m:
            return LocTuple('rbp', int(m.group(1)))

        m = re.search(r'DW_OP_reg\d+ \((\w+)\)', line)
        if m:
            return LocTuple(m.group(1).lower(), None)

        m = re.search(r'DW_OP_addr: (-?[\da-fA-F]+)', line)
        if m:
            return LocTuple(None, int(m.group(1), 16))

        m = re.search(r'DW_OP_breg\d+ \(([a-zA-Z]+)\): (-?\d+)', line)
        if m:
            return LocTuple(m.group(1).lower(), int(m.group(2)))

        m = re.search(r'DW_OP_breg\d+ \([a-zA-Z]+\)\(-?\d+\)', line)
        if m:
            return LocTuple(m.group(1).lower(), int(m.group(2)))

        return LocTuple(None, None)

    def __bool__(self):
        return not (self.register is None and self.offset is None)

    def __repr__(self):
        if self.register is None:
            if self.offset is None:
                return '<None>'
            else:
                return '0x%x' % self.offset
        else:
            if self.offset is None:
                return self.register
            else:
                return f'({self.register}){self.offset}'

    def is_static(self):
        return self.register is None and self.offset is not None


class DebinStats:
    def __init__(self, truth_path: str, inferred_path: str):
        truth_file = open(truth_path, 'rb')
        truth_elf_file = ELFFile(truth_file)
        if not truth_elf_file.has_dwarf_info():
            print(f'"{truth_path}" has no DWARF info')
            return

        inferred_file = open(inferred_path, 'rb')
        inferred_elf_file = ELFFile(inferred_file)
        if not inferred_elf_file.has_dwarf_info():
            print(f'"{inferred_path}" has no DWARF info')
            return

        assert truth_elf_file.get_machine_arch() == inferred_elf_file.get_machine_arch()
        set_global_machine_arch(truth_elf_file.get_machine_arch())

        self.truth_dwarf_info = truth_elf_file.get_dwarf_info()
        self.inferred_dwarf_info = inferred_elf_file.get_dwarf_info()

        # For parsing ExprLoc
        self.truth_visitor = ExprDumper(self.truth_dwarf_info.structs)
        self.inferred_visitor = ExprDumper(self.inferred_dwarf_info.structs)

        self.num_truth = 0
        self.num_tried = 0
        self.num_correct = 0

        self.inferred_funcs_die: Dict[int, Optional[Tuple[CompileUnit, DIE]]] = dict()
        self.inferred_static_var_types: Dict[int, str] = dict()
        self.inferred_funcs_var_types: Dict[int, Optional[Dict[LocTuple, str]]] = dict()

    def run(self):
        for cu in self.truth_dwarf_info.iter_CUs():
            log('Found a CU at offset 0x%x, die_offset 0x%x, length 0x%x %s' %
                (cu.cu_offset, cu.cu_die_offset, cu['unit_length'], '*' * 80))

            top_die: DIE = cu.get_top_DIE()
            self.stat_die_rec(top_die, None, cu)

    def stat_die_rec(self, die: DIE, parent_fn: Optional[DIE], cu: CompileUnit):
        if die.tag is None:
            return
        if die.tag in ['DW_TAG_variable', 'DW_TAG_formal_parameter']:
            return self.stat_variable(die, parent_fn, cu)
        if die.tag == 'DW_TAG_subprogram':
            for child in die.iter_children():
                self.stat_die_rec(child, die, cu)
            return

        for child in die.iter_children():
            self.stat_die_rec(child, parent_fn, cu)

    def stat_variable(self, var_die: DIE, parent_fn: Optional[DIE], cu: CompileUnit):
        assert var_die.tag in ['DW_TAG_variable', 'DW_TAG_formal_parameter']
        # Get correct type
        real_type_str: str = DebinStats.get_type_str(var_die, cu)
        # log(f'{die_name(var_die)}: {real_type_str}')
        if not real_type_str:
            return
        # Filter non STL variables
        if not(any(real_type_str.startswith(s+'<') for s in ['vector', 'list', 'map',
                                                             'std::vector', 'std::list', 'std::map'])):
            return
        # Get location
        location_str = DebinStats.get_location_str(var_die, self.truth_visitor)
        if location_str is None:
            return
        # Get locTuple
        loc_tuple = LocTuple.from_str(location_str)
        if not loc_tuple:
            return

        log(f'\nstat_variable({die_name(var_die)}: {real_type_str}, {die_name(parent_fn)})\n{location_str} => {loc_tuple}')
        self.num_truth += 1
        cu_die = self.search_matching(loc_tuple, parent_fn)
        if cu_die is None:
            return
        else:
            self.num_tried += 1
            self.type_match(real_type_str, cu_die[0], cu_die[1])

    def search_matching(self, target_loc: LocTuple, parent_fn: Optional[DIE]) -> Optional[Tuple[CompileUnit, DIE]]:
        if target_loc.is_static():
            return self.search_matching_static(target_loc)
        else:
            return self.search_matching_auto(target_loc, parent_fn)

    def search_matching_static(self, target_loc_tuple: LocTuple) -> Optional[Tuple[CompileUnit, DIE]]:
        for cu in self.inferred_dwarf_info.iter_CUs():
            for die in cu.iter_DIEs():  # type: DIE
                if die.tag not in ['DW_TAG_variable', 'DW_TAG_formal_parameter']:
                    continue
                location_str = DebinStats.get_location_str(die, self.inferred_visitor)
                if location_str is None:
                    continue
                loc_tuple = LocTuple.from_str(location_str)
                if not loc_tuple:
                    continue
                if loc_tuple == target_loc_tuple:
                    return cu, die
        return None

    def get_matching_func_die(self, low_pc: int) -> Optional[Tuple[CompileUnit, DIE]]:
        key_low_pc = 'DW_AT_low_pc'

        # None means no match, 0 means haven't tried
        cu_die = self.inferred_funcs_die.get(low_pc, 0)
        if cu_die != 0:
            return cu_die

        for cu in self.inferred_dwarf_info.iter_CUs():
            for die in cu.iter_DIEs():  # type: DIE
                if die.tag != 'DW_TAG_subprogram':
                    continue
                fn_low_pc = die.attributes.get(key_low_pc).value
                if fn_low_pc is None:
                    continue
                if fn_low_pc != low_pc:
                    continue
                # Now we can assure that the 'die' is our target function
                result = (cu, die)
                self.inferred_funcs_die[low_pc] = result
                return result
        # Couldn't find one
        self.inferred_funcs_die[low_pc] = None
        return None

    def search_matching_auto(self, target_loc_tuple: LocTuple, parent_fn: DIE) -> Optional[Tuple[CompileUnit, DIE]]:
        key_low_pc = 'DW_AT_low_pc'
        parent_low_pc = parent_fn.attributes.get(key_low_pc).value
        if parent_low_pc is None:
            return
        log('My low pc is %08x' % parent_low_pc)

        cu_die = self.get_matching_func_die(parent_low_pc)
        if cu_die is None:
            return
        cu, die = cu_die

        # Closure
        def search_within_die(entry: DIE) -> Optional[DIE]:
            if entry.tag in ['DW_TAG_variable', 'DW_TAG_formal_parameter']:
                location_str = DebinStats.get_location_str(entry, self.inferred_visitor)
                if location_str is None:
                    return
                if LocTuple.from_str(location_str) == target_loc_tuple:
                    return entry
            else:
                for child in entry.iter_children():     # type: DIE
                    if child.tag is None or child.tag == 'DW_TAG_subprogram':
                        continue
                    result = search_within_die(child)
                    if result is not None:
                        return result

        searched_child = search_within_die(die)
        if searched_child is None:
            log('\t\t\tNot found')
            return None
        log(f'\t\t\tFound! type:"{DebinStats.get_type_str(searched_child, cu)}", offset={searched_child.offset}')
        return cu, searched_child

    def type_match(self, correct_type_str: str, inferred_cu: CompileUnit, inferred_die: DIE):
        log(f'type_match({correct_type_str}, {inferred_cu}, {inferred_die})')
        # Get type of inferred DIE
        inferred_type_str: str = DebinStats.get_type_str(inferred_die, inferred_cu)
        log(f'inferred_type_str = {inferred_type_str}')
        if not inferred_type_str:
            # TODO: know variable but has no type
            return
        # print(f'Inferred: {inferred_type_str}\t<- Truth: {correct_type_str}')
        if (correct_type_str == inferred_type_str
                or correct_type_str.endswith('*') and inferred_type_str.endswith('*')
                or ('int' in correct_type_str and 'int' in inferred_type_str)
                or ('long' in correct_type_str and 'long' in inferred_type_str)
                or ('array' in correct_type_str and 'array' in inferred_type_str)
                or (not correct_type_str.endswith('*')) and inferred_type_str.lower() == 'struct'):
            #print('\t\tRight!!!')
            self.num_correct += 1   # TODO: ...
        else:
            pass#print('\t\tWrong...')

    @staticmethod
    def get_location_str(die: DIE, visitor: ExprDumper) -> Optional[str]:
        visitor.clear()
        location: AttributeValue = die.attributes.get('DW_AT_location')
        if location is None:
            return None
        try:
            visitor.process_expr(location.value)
        except TypeError:
            return None
        return visitor.get_str()

    @staticmethod
    def get_type_str(die: DIE, cu: CompileUnit) -> str:
        return DebinStats.base_type_to_str(die, cu)

    @staticmethod
    def base_type_to_str(entry: DIE, cu: CompileUnit) -> str:
        tag: str = entry.tag
        # if tag in ['DW_TAG_structure_type', 'DW_TAG_class_type', 'DW_TAG_base_type']:
        #     return entry.attributes.get('name', '{NoName}')

        base_type: AttributeValue = entry.attributes.get('DW_AT_type')
        if base_type is None:
            #log(f'{entry.offset} no base_type')
            name_attr: AttributeValue = entry.attributes.get('DW_AT_name')
            if name_attr is not None:
                return name_attr.value.decode('utf-8')
            else:
                return {
                    'DW_TAG_array_type':            'array',
                    'DW_TAG_enumeration_type':      'ENUM',
                    'DW_TAG_pointer_type':          '*',
                    'DW_TAG_reference_type':        '*',
                    'DW_TAG_rvalue_reference_type': '*',
                    'DW_TAG_structure_type':        'STRUCT',
                    'DW_TAG_class_type':            'STRUCT',
                }.get(tag, '')

        #log(f'{entry.offset} HAVE base_type')
        base_entry = DebinStats.get_die_in_cu_by_offset(cu, base_type.value)

        fmt_str_dict = {
            'DW_TAG_const_type':            '{}',
            'DW_TAG_typedef':               '{}',
            'DW_TAG_variable':              '{}',
            'DW_TAG_rvalue_reference_type': '{}',
            'DW_TAG_formal_parameter':      '{}',
            'DW_TAG_volatile_type':         '{}',
            'DW_TAG_array_type':            'array{}',
            'DW_TAG_pointer_type':          '{}*',
            'DW_TAG_reference_type':        '{}*',
            'DW_TAG_enumeration_type':      'ENUM',
        }
        fmt_str = fmt_str_dict.get(tag)
        if fmt_str is not None:
            if '{}' in fmt_str:
                return fmt_str.format(DebinStats.base_type_to_str(base_entry, cu))
            else:
                return fmt_str

        return ''   # tag

    @staticmethod
    def get_die_in_cu_by_offset(cu: CompileUnit, offset) -> Optional[DIE]:
        left = 0
        right = len(cu._dielist) - 1
        real_offset = cu.cu_offset + offset

        while left <= right:
            mid = left + (right - left) // 2
            mid_entry: DIE = cu._get_DIE(mid)
            if mid_entry.offset < real_offset:
                left = mid + 1
            elif mid_entry.offset > real_offset:
                right = mid - 1
            else:
                return mid_entry
        return None


if __name__ == '__main__':
    if len(sys.argv) < 3:
        print(f'Usage: {sys.argv[0]} <debug-binary> <inferred-binary>')
    if '-c' in sys.argv:
        cache_enabled = True

    stats = DebinStats(sys.argv[1], sys.argv[2])

    print(f'Comparing "{sys.argv[1]}" against "{sys.argv[2]}"...')
    tic = datetime.now()
    stats.run()
    toc = datetime.now()
    time_diff = toc - tic

    def percentage(m, n):
        if not n:
            return 'NaN'
        return '{:.2f}%'.format(100 * m / n)


    print(
        f'N(correct) = {stats.num_correct}',
        f'  N(tried)   = {stats.num_tried}',
        f'    correct / tried = {percentage(stats.num_correct, stats.num_tried)}',
        f'  N(truth)   = {stats.num_truth}',
        f'    correct / truth = {percentage(stats.num_correct, stats.num_truth)}',
        'Total time: {:.3f} second(s)'.format(time_diff.total_seconds()),
        sep='\n'
    )

